package com.leetcode.No1617;

import org.junit.Test;

import java.util.*;

/**
 * 1617: 统计子树中城市之间最大距离
 */
public class Solution {

    int[][] edges;
    boolean[] visited;
    Map<Integer, Node> nodeMap;

    public int[] countSubgraphsForEachDiameter(int n, int[][] edges) {
        this.edges = edges;
        visited = new boolean[n];
        nodeMap = new HashMap<>();

        initTree(1);
        calDistance(1);

        int[] result = new int[n - 1];
        for (Node node : nodeMap.values()) {
            int idx = node.maxDis - 1;
            if (idx >= 0) {
                ++result[idx];
            }
//            node.
        }
        return result;
    }

    private void initTree(int nodeId) {
        Node node = new Node(nodeId);
        visited[nodeId - 1] = true;
        nodeMap.put(nodeId, node);

        // 找到子节点
        for (int[] edge : edges) {
            int other;
            if (edge[0] == nodeId) {
                other = edge[1];
            } else if (edge[1] == nodeId) {
                other = edge[0];
            } else {
                continue;
            }
            if (!visited[other - 1]) {
                node.sonNode.add(other);
                initTree(other);
            }
        }

    }

    private int calDistance(int i) {
        Node fatherNode = nodeMap.get(i);
        List<Integer> sonNode = fatherNode.sonNode;
        List<Integer> nodeDis = fatherNode.nodeDis;

        if (sonNode.size() == 0) {
            return 0;
        }

        int firstDis = 0;
        int secondDis = 0;

        for (Integer sonNodeId : sonNode) {
            int dis = 1 + calDistance(sonNodeId);
            nodeDis.add(dis);
            if (firstDis <= dis) {
                secondDis = firstDis;
                firstDis = dis;
            } else if (secondDis < dis) {
                secondDis = dis;
            }
        }

        int maxDis = firstDis + secondDis;
        for (Integer sonNodeId : sonNode) {
            Node node = nodeMap.get(sonNodeId);
            maxDis = Math.max(maxDis, node.maxDis);
        }
        fatherNode.maxDis = maxDis;

        return firstDis;
    }

    // 除了叶子结点,每个node都会引出以它自己为root的一棵子树,
    // 可以认为每个非叶子node各会引出一棵互不相同的子树
    private class Node {
        int id;
        // 本节点的直接子节点
        List<Integer> sonNode = new ArrayList<>();
        // 本节点 到 以本节点为root的子树中的所有叶子节点 的距离
        List<Integer> nodeDis = new ArrayList<>();
        // 以本节点为root的子树中 的 两节点间最远距离,
        // 这个距离可能是父结点两个子树之间的距离,也可能是子树的两个子树之间的距离
        int maxDis = 0;

        public Node(int id) {
            this.id = id;
        }
    }

    @Test
    public void test1() {
        int n;
        int[][] edges;

        n = 4;
        edges = new int[][]{{1,2},{2,3},{2,4}};
        System.out.println(Arrays.toString(countSubgraphsForEachDiameter(n, edges)));
    }
}
